import os
import yaml
from mindspore.ops import operations as P
from mindspore import jit_class
from .dump_base import DumpBase
from ..common.utils import Const
from troubleshooter.migrator.save_for_checker import save_tensor


class PrimitiveGenerator:
    def __init__(self):
        self.support_dict = {}
        self.gen_dict = {}

    @staticmethod
    def new_init(self, *args, **kwargs):
        DumpBase.__init__(self)
        if self.__class__.args_cache:
            if [args, kwargs] != self.__class__.args_cache:
                self.__class__.args_cache = [args, kwargs]
                self.__class__.global_index += 1
        else:
            self.__class__.args_cache = [args, kwargs]
        self.op = self.__class__.origin_func(*args, **kwargs)
        self.instance_index = self.__class__.global_index
        self.set_new_func(self.__class__, self.instance_index, self.__class__.origin_func(*args, **kwargs))
        self.save_path = self.get_save_path(self.__class__.api_name, self.instance_index)
        self.info_dict = self.save_info(args, kwargs, self.save_path)
        self.__class__.global_dict[self.__class__.api_name + "_" + str(self.instance_index)] = self.info_dict

    @staticmethod
    def new_call(self, *args):
        save_tensor(
            self.save_path + "/" + self.__class__.api_name + "_" + str(self.instance_index) + "_" + Const.INPUT,
            args
        )
        instance = self.get_func(self.__class__, str(self.instance_index))
        out = instance(*args)
        save_tensor(
            self.save_path + "/" + self.__class__.api_name + "_" + str(self.instance_index) + "_" + Const.OUTPUT,
            out
        )

        return out

    def get_support_ops(self):
        cur_path = os.path.dirname(os.path.realpath(__file__))
        yaml_path = os.path.join(cur_path, "support_ops.yaml")
        with open(yaml_path, "r") as f:
            support_ops = yaml.safe_load(f).get("Primitive")
        
        for ops in dir(P):
            if ops in support_ops:
                self.support_dict[ops] = getattr(P, ops)

    def gen_ops(self):
        for name, ops in self.support_dict.items():
            gen_cls = type(
                "Dump" + name,
                (DumpBase,),
                {
                    "api_name": ops.__name__,
                    "global_index": 0,
                    "global_dict": {},
                    "origin_func": ops,
                    "args_cache": [],
                    "__init__":self.new_init,
                    "__call__":self.new_call,
                }
            )
        
            gen_cls = jit_class(gen_cls)
            self.gen_dict[name] = gen_cls

    def init_dump(self, cls):
        self.get_support_ops()
        self.gen_ops()
        for name, ops in self.gen_dict.items():
            setattr(cls, name, ops)


def wrap_cls(cls):
    gen = PrimitiveGenerator()
    gen.init_dump(cls)
    return cls


class Meta(type):
    gen = PrimitiveGenerator()
    gen.get_support_ops()
    support_dict = gen.support_dict
    def __getattr__(cls, name):
        if name not in Meta.support_dict:
            return getattr(P, name)
        else:
            return getattr(cls, name)


@wrap_cls
class Q(metaclass=Meta):
    pass
